#!/bin/bash

# Jetson Nano Memory Optimization Script
# 保留Dify AI服务的内存优化脚本

echo "=== Jetson Nano Memory Optimization Script ==="
echo "开始内存优化，保留Dify AI服务..."

# 备份当前配置
echo "备份当前系统配置..."
sudo cp /etc/sysctl.conf /etc/sysctl.conf.backup.$(date +%Y%m%d_%H%M%S)

# 1. 优化内存管理参数
echo "优化内存管理参数..."

# 创建内存优化配置文件
sudo tee /etc/sysctl.d/99-memory-optimization.conf > /dev/null << 'EOF'
# Memory optimization for Jetson Nano
# 减少swap使用，优先使用物理内存
vm.swappiness = 1

# 增加缓存压力，更积极地回收缓存
vm.vfs_cache_pressure = 80

# 优化脏页处理
vm.dirty_ratio = 5
vm.dirty_background_ratio = 2

# 减少内存碎片
vm.min_free_kbytes = 65536

# 优化网络缓冲区
net.core.rmem_max = 134217728
net.core.wmem_max = 134217728
net.core.netdev_max_backlog = 5000

# 减少内核日志缓冲区
kernel.printk = 3 4 1 3
EOF

# 应用内存优化参数
echo "应用内存优化参数..."
sudo sysctl -p /etc/sysctl.d/99-memory-optimization.conf

# 2. 禁用不必要的服务（保留Dify AI相关服务）
echo "禁用不必要的系统服务..."

# 定义可以安全禁用的服务列表（保留USB挂载功能）
SERVICES_TO_DISABLE=(
    "snapd.service"
    "snapd.socket"
    "snapd.seeded.service"
    "snapd.autoimport.service"
    "snapd.core-fixup.service"
    "snapd.recovery-chooser-trigger.service"
    "snapd.system-shutdown.service"
    "snapd.apparmor.service"
    "snapd.snap-repair.timer"
    "ModemManager.service"
    "spice-vdagent.service"
    "spice-vdagentd.service"
    "kerneloops.service"        # 禁用内核错误报告
    "anacron.service"           # 禁用定时任务功能
    "anacron.timer"             # 禁用定时任务功能
    "accounts-daemon.service"   # 禁用用户账户管理
    "whoopsie.service"          # 禁用错误报告
    "apport.service"            # 禁用应用程序崩溃报告
    "apport-autoreport.path"    # 禁用自动错误报告
    "apport-forward.socket"     # 禁用错误报告转发
    # 注释掉以下服务以保留功能：
    # "udisks2.service"           # 保留USB设备自动挂载
)

# 禁用服务
for service in "${SERVICES_TO_DISABLE[@]}"; do
    if systemctl is-enabled "$service" >/dev/null 2>&1; then
        echo "禁用服务: $service"
        sudo systemctl disable "$service" >/dev/null 2>&1
        sudo systemctl stop "$service" >/dev/null 2>&1
    else
        echo "服务 $service 已经被禁用或不存在"
    fi
done

# 3. 清理系统缓存和临时文件
echo "清理系统缓存..."
sudo sync
echo 3 | sudo tee /proc/sys/vm/drop_caches > /dev/null
sudo apt-get clean
sudo apt-get autoremove -y

# 4. 优化Docker（如果正在使用）
if systemctl is-active docker >/dev/null 2>&1; then
    echo "优化Docker配置..."
    
    # 创建Docker优化配置
    sudo mkdir -p /etc/docker
    sudo tee /etc/docker/daemon.json > /dev/null << 'EOF'
{
    "log-driver": "json-file",
    "log-opts": {
        "max-size": "10m",
        "max-file": "3"
    },
    "storage-driver": "overlay2",
    "default-ulimits": {
        "memlock": {
            "name": "memlock",
            "soft": -1,
            "hard": -1
        }
    }
}
EOF
    
    echo "重启Docker服务..."
    sudo systemctl restart docker
fi

# 5. 创建内存监控脚本
echo "创建内存监控脚本..."
sudo tee /usr/local/bin/memory-monitor.sh > /dev/null << 'EOF'
#!/bin/bash
# 内存监控脚本

echo "=== 内存使用情况 $(date) ==="
free -h
echo ""
echo "=== 占用内存最多的进程 ==="
ps aux --sort=-%mem | head -10
echo ""
echo "=== Swap使用情况 ==="
swapon --show
echo ""
echo "=== 系统负载 ==="
uptime
EOF

sudo chmod +x /usr/local/bin/memory-monitor.sh

# 6. 验证Dify AI服务状态
echo "验证Dify AI相关服务状态..."
echo "检查Celery服务..."
pgrep -f celery && echo "Celery服务正在运行" || echo "Celery服务未运行"

echo "检查Gunicorn服务..."
pgrep -f gunicorn && echo "Gunicorn服务正在运行" || echo "Gunicorn服务未运行"

echo "检查Docker容器..."
if command -v docker >/dev/null 2>&1; then
    docker ps --format "table {{.Names}}\t{{.Status}}\t{{.Ports}}" | grep -E "(dify|postgres|redis|nginx)"
fi

# 7. 显示优化结果
echo ""
echo "=== 内存优化完成 ==="
echo "当前内存使用情况："
free -h
echo ""
echo "已禁用的服务数量: ${#SERVICES_TO_DISABLE[@]}"
echo "保留的功能："
echo "✅ USB设备自动挂载 (udisks2.service)"
echo ""
echo "已禁用的功能（释放更多内存）："
echo "❌ 定时任务调度 (anacron.service)"
echo "❌ 错误自动报告 (apport, whoopsie, kerneloops)"
echo "❌ 用户账户管理 (accounts-daemon.service)"
echo ""
echo "优化建议："
echo "1. 运行 'memory-monitor.sh' 监控内存使用"
echo "2. 重启系统以确保所有优化生效"
echo "3. 监控Dify AI服务是否正常运行"
echo ""
echo "如需回滚，请运行："
echo "sudo cp /etc/sysctl.conf.backup.* /etc/sysctl.conf"
echo "sudo systemctl enable [服务名] # 重新启用被禁用的服务"